Tiny-cuda-nn 环境配置

Tiny-cuda-nn，一个轻量级的神经网络算子加速库，可用于 Pytorch 扩展

研究其代码后，发现其主要依靠 Pytorch 内置的 Cpp_extension 功能来编译为用户扩展算子。相关 Pytorch 特性可以查看下方链接

Pytorch 自定义 C++ 和 CUDA 扩展

下文介绍如何配置 Tiny-cuda-nn 环境

环境配置

参照 Tiny-cuda-nn 环境配置官方文档

  1. 检查 CUDA 环境

首先检查全局 CUDA 环境是否安装，一般查看是否存在路径 /usr/local/cuda 即可

其次检查全局的 nvcc 编译器版本是否匹配，通过命令 nvcc -version 查看 nvcc 版本，与 CUDA 版本进行比对，确保大版本号一致

  1. 准备 Python 环境

依据官方文档，使用 Pytorch 的 2.0.0 版本，Python 版本无限制，此处以 Python 3.11 版本演示

1
2
3
conda create -n env_name python=3.11
conda activate env_name
pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1 --index-url https://download.pytorch.org/whl/cu118
  1. 编译 tiny-cuda-nn

接着借助 Pytorch 的 Cpp 扩展来编译 Tiny-cuda-nn，此部分大概用时 30min

1
2
3
git clone --recursive https://github.com/nvlabs/tiny-cuda-nn
cd tiny-cuda-nn/bindings/torch
python setup.py install
  1. 检查

不出意外，一个叫做 tinycudann 的 pip 第三方包就安装完毕了

运行下方代码检查是否安装成功

1
2
3
import tinycudann

print(tinycudann.__all__)

可能报错及原因分析

  1. Pytorch 无法使用 CUDA

请检查 Pytorch 下载的 cuda 版本是否与主机 cuda 版本匹配

  1. 找不到 CUDA_HOME

请检查全局 CUDA 环境是否安装，一般查看是否存在路径 /usr/local/cuda 即可

  1. 编译器 CUDA 版本不匹配

请检查全局的 nvcc 编译器版本是否匹配，通过命令 nvcc -version 查看 nvcc 版本，与 CUDA 版本进行比对，确保大版本号一致

  1. 编译时发现 libtorch 头文件缺失

请检查 Pytorch 版本是否为 2.0.0，若版本过高会有头文件缺失的情况

后言

墨子号说配这个环境搞了一周也没搞定，于是我今晚看了一下，发现涉及到很多 cuda 底层版本对应和编译的东西，着实头疼

两三个小时配了一下，并不复杂，顺便把各种报错及其原因也分析了一遍，收获很多

期间也是览了很多 Pytorch 的算子扩展原理、CUDA 编译原理之类的，确实很有趣